{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67,
     "referenced_widgets": [
      "a8e97e26d65147efa2ddb176214e818a",
      "19e913a07b044b2d91ca7b5d2dbbe8f7",
      "1de97479fc3e44caae336cf0b1082d02",
      "b8d4e152c232416f82f5fb30e2478ef8",
      "d66b3f8971b24c81abaefe111e6383a6",
      "11d14b9bf0f74bf7bd1b58e6a26fd49e",
      "e9fde27bd58b499bb386b36314d9efca",
      "5407db80343c44c79350021cd98204f8",
      "68e9f5994766480dac5dc77ce043a2e9",
      "d00c195ac5644bd0aba81dccf56e2df0",
      "54c35b52c4e24543872c8335c935c410"
     ]
    },
    "id": "VrX4VTl5pYNq",
    "outputId": "63f38652-f298-4723-bd5c-e1bbe5f9d961"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Imports**\n",
    "# diffusion policy import\n",
    "from typing import Tuple, Sequence, Dict, Union, Optional, Callable\n",
    "import numpy as np\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import collections\n",
    "import zarr\n",
    "from diffusers.schedulers.scheduling_ddpm import DDPMScheduler\n",
    "from diffusers.training_utils import EMAModel\n",
    "from diffusers.optimization import get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# env import\n",
    "import gym\n",
    "from gym import spaces\n",
    "import pygame\n",
    "import pymunk\n",
    "import pymunk.pygame_util\n",
    "from pymunk.space_debug_draw_options import SpaceDebugColor\n",
    "from pymunk.vec2d import Vec2d\n",
    "import shapely.geometry as sg\n",
    "import cv2\n",
    "import skimage.transform as st\n",
    "from skvideo.io import vwrite\n",
    "from IPython.display import Video\n",
    "import gdown\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "L5E-nR6ornyg"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Environment**\n",
    "#@markdown Defines a PyMunk-based Push-T environment `PushTEnv`.\n",
    "#@markdown And it's subclass `PushTImageEnv`.\n",
    "#@markdown\n",
    "#@markdown **Goal**: push the gray T-block into the green area.\n",
    "#@markdown\n",
    "#@markdown Adapted from [Implicit Behavior Cloning](https://implicitbc.github.io/)\n",
    "\n",
    "\n",
    "positive_y_is_up: bool = False\n",
    "\"\"\"Make increasing values of y point upwards.\n",
    "\n",
    "When True::\n",
    "\n",
    "    y\n",
    "    ^\n",
    "    |      . (3, 3)\n",
    "    |\n",
    "    |   . (2, 2)\n",
    "    |\n",
    "    +------ > x\n",
    "\n",
    "When False::\n",
    "\n",
    "    +------ > x\n",
    "    |\n",
    "    |   . (2, 2)\n",
    "    |\n",
    "    |      . (3, 3)\n",
    "    v\n",
    "    y\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:\n",
    "    \"\"\"Convenience method to convert pymunk coordinates to pygame surface\n",
    "    local coordinates.\n",
    "\n",
    "    Note that in case positive_y_is_up is False, this function wont actually do\n",
    "    anything except converting the point to integers.\n",
    "    \"\"\"\n",
    "    if positive_y_is_up:\n",
    "        return round(p[0]), surface.get_height() - round(p[1])\n",
    "    else:\n",
    "        return round(p[0]), round(p[1])\n",
    "\n",
    "\n",
    "def light_color(color: SpaceDebugColor):\n",
    "    color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255]))\n",
    "    color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3])\n",
    "    return color\n",
    "\n",
    "class DrawOptions(pymunk.SpaceDebugDrawOptions):\n",
    "    def __init__(self, surface: pygame.Surface) -> None:\n",
    "        \"\"\"Draw a pymunk.Space on a pygame.Surface object.\n",
    "\n",
    "        Typical usage::\n",
    "\n",
    "        >>> import pymunk\n",
    "        >>> surface = pygame.Surface((10,10))\n",
    "        >>> space = pymunk.Space()\n",
    "        >>> options = pymunk.pygame_util.DrawOptions(surface)\n",
    "        >>> space.debug_draw(options)\n",
    "\n",
    "        You can control the color of a shape by setting shape.color to the color\n",
    "        you want it drawn in::\n",
    "\n",
    "        >>> c = pymunk.Circle(None, 10)\n",
    "        >>> c.color = pygame.Color(\"pink\")\n",
    "\n",
    "        See pygame_util.demo.py for a full example\n",
    "\n",
    "        Since pygame uses a coordiante system where y points down (in contrast\n",
    "        to many other cases), you either have to make the physics simulation\n",
    "        with Pymunk also behave in that way, or flip everything when you draw.\n",
    "\n",
    "        The easiest is probably to just make the simulation behave the same\n",
    "        way as Pygame does. In that way all coordinates used are in the same\n",
    "        orientation and easy to reason about::\n",
    "\n",
    "        >>> space = pymunk.Space()\n",
    "        >>> space.gravity = (0, -1000)\n",
    "        >>> body = pymunk.Body()\n",
    "        >>> body.position = (0, 0) # will be positioned in the top left corner\n",
    "        >>> space.debug_draw(options)\n",
    "\n",
    "        To flip the drawing its possible to set the module property\n",
    "        :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip\n",
    "        the simulation upside down before drawing::\n",
    "\n",
    "        >>> positive_y_is_up = True\n",
    "        >>> body = pymunk.Body()\n",
    "        >>> body.position = (0, 0)\n",
    "        >>> # Body will be position in bottom left corner\n",
    "\n",
    "        :Parameters:\n",
    "                surface : pygame.Surface\n",
    "                    Surface that the objects will be drawn on\n",
    "        \"\"\"\n",
    "        self.surface = surface\n",
    "        super(DrawOptions, self).__init__()\n",
    "\n",
    "    def draw_circle(\n",
    "        self,\n",
    "        pos: Vec2d,\n",
    "        angle: float,\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        p = to_pygame(pos, self.surface)\n",
    "\n",
    "        pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0)\n",
    "        pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius-4), 0)\n",
    "\n",
    "        circle_edge = pos + Vec2d(radius, 0).rotated(angle)\n",
    "        p2 = to_pygame(circle_edge, self.surface)\n",
    "        line_r = 2 if radius > 20 else 1\n",
    "        # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r)\n",
    "\n",
    "    def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None:\n",
    "        p1 = to_pygame(a, self.surface)\n",
    "        p2 = to_pygame(b, self.surface)\n",
    "\n",
    "        pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2])\n",
    "\n",
    "    def draw_fat_segment(\n",
    "        self,\n",
    "        a: Tuple[float, float],\n",
    "        b: Tuple[float, float],\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        p1 = to_pygame(a, self.surface)\n",
    "        p2 = to_pygame(b, self.surface)\n",
    "\n",
    "        r = round(max(1, radius * 2))\n",
    "        pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r)\n",
    "        if r > 2:\n",
    "            orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])]\n",
    "            if orthog[0] == 0 and orthog[1] == 0:\n",
    "                return\n",
    "            scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5\n",
    "            orthog[0] = round(orthog[0] * scale)\n",
    "            orthog[1] = round(orthog[1] * scale)\n",
    "            points = [\n",
    "                (p1[0] - orthog[0], p1[1] - orthog[1]),\n",
    "                (p1[0] + orthog[0], p1[1] + orthog[1]),\n",
    "                (p2[0] + orthog[0], p2[1] + orthog[1]),\n",
    "                (p2[0] - orthog[0], p2[1] - orthog[1]),\n",
    "            ]\n",
    "            pygame.draw.polygon(self.surface, fill_color.as_int(), points)\n",
    "            pygame.draw.circle(\n",
    "                self.surface,\n",
    "                fill_color.as_int(),\n",
    "                (round(p1[0]), round(p1[1])),\n",
    "                round(radius),\n",
    "            )\n",
    "            pygame.draw.circle(\n",
    "                self.surface,\n",
    "                fill_color.as_int(),\n",
    "                (round(p2[0]), round(p2[1])),\n",
    "                round(radius),\n",
    "            )\n",
    "\n",
    "    def draw_polygon(\n",
    "        self,\n",
    "        verts: Sequence[Tuple[float, float]],\n",
    "        radius: float,\n",
    "        outline_color: SpaceDebugColor,\n",
    "        fill_color: SpaceDebugColor,\n",
    "    ) -> None:\n",
    "        ps = [to_pygame(v, self.surface) for v in verts]\n",
    "        ps += [ps[0]]\n",
    "\n",
    "        radius = 2\n",
    "        pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps)\n",
    "\n",
    "        if radius > 0:\n",
    "            for i in range(len(verts)):\n",
    "                a = verts[i]\n",
    "                b = verts[(i + 1) % len(verts)]\n",
    "                self.draw_fat_segment(a, b, radius, fill_color, fill_color)\n",
    "\n",
    "    def draw_dot(\n",
    "        self, size: float, pos: Tuple[float, float], color: SpaceDebugColor\n",
    "    ) -> None:\n",
    "        p = to_pygame(pos, self.surface)\n",
    "        pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0)\n",
    "\n",
    "\n",
    "def pymunk_to_shapely(body, shapes):\n",
    "    geoms = list()\n",
    "    for shape in shapes:\n",
    "        if isinstance(shape, pymunk.shapes.Poly):\n",
    "            verts = [body.local_to_world(v) for v in shape.get_vertices()]\n",
    "            verts += [verts[0]]\n",
    "            geoms.append(sg.Polygon(verts))\n",
    "        else:\n",
    "            raise RuntimeError(f'Unsupported shape type {type(shape)}')\n",
    "    geom = sg.MultiPolygon(geoms)\n",
    "    return geom\n",
    "\n",
    "# env\n",
    "class PushTEnv(gym.Env):\n",
    "    metadata = {\"render.modes\": [\"human\", \"rgb_array\"], \"video.frames_per_second\": 10}\n",
    "    reward_range = (0., 1.)\n",
    "\n",
    "    def __init__(self,\n",
    "            legacy=False,\n",
    "            block_cog=None, damping=None,\n",
    "            render_action=True,\n",
    "            render_size=96,\n",
    "            reset_to_state=None\n",
    "        ):\n",
    "        self._seed = None\n",
    "        self.seed()\n",
    "        self.window_size = ws = 512  # The size of the PyGame window\n",
    "        self.render_size = render_size\n",
    "        self.sim_hz = 100\n",
    "        # Local controller params.\n",
    "        self.k_p, self.k_v = 100, 20    # PD control.z\n",
    "        self.control_hz = self.metadata['video.frames_per_second']\n",
    "        # legcay set_state for data compatiblity\n",
    "        self.legacy = legacy\n",
    "\n",
    "        # agent_pos, block_pos, block_angle\n",
    "        self.observation_space = spaces.Box(\n",
    "            low=np.array([0,0,0,0,0], dtype=np.float64),\n",
    "            high=np.array([ws,ws,ws,ws,np.pi*2], dtype=np.float64),\n",
    "            shape=(5,),\n",
    "            dtype=np.float64\n",
    "        )\n",
    "\n",
    "        # positional goal for agent\n",
    "        self.action_space = spaces.Box(\n",
    "            low=np.array([0,0], dtype=np.float64),\n",
    "            high=np.array([ws,ws], dtype=np.float64),\n",
    "            shape=(2,),\n",
    "            dtype=np.float64\n",
    "        )\n",
    "\n",
    "        self.block_cog = block_cog\n",
    "        self.damping = damping\n",
    "        self.render_action = render_action\n",
    "\n",
    "        \"\"\"\n",
    "        If human-rendering is used, `self.window` will be a reference\n",
    "        to the window that we draw to. `self.clock` will be a clock that is used\n",
    "        to ensure that the environment is rendered at the correct framerate in\n",
    "        human-mode. They will remain `None` until human-mode is used for the\n",
    "        first time.\n",
    "        \"\"\"\n",
    "        self.window = None\n",
    "        self.clock = None\n",
    "        self.screen = None\n",
    "\n",
    "        self.space = None\n",
    "        self.teleop = None\n",
    "        self.render_buffer = None\n",
    "        self.latest_action = None\n",
    "        self.reset_to_state = reset_to_state\n",
    "\n",
    "    def reset(self):\n",
    "        seed = self._seed\n",
    "        self._setup()\n",
    "        if self.block_cog is not None:\n",
    "            self.block.center_of_gravity = self.block_cog\n",
    "        if self.damping is not None:\n",
    "            self.space.damping = self.damping\n",
    "\n",
    "        # use legacy RandomState for compatiblity\n",
    "        state = self.reset_to_state\n",
    "        if state is None:\n",
    "            rs = np.random.RandomState(seed=seed)\n",
    "            state = np.array([\n",
    "                rs.randint(50, 450), rs.randint(50, 450),\n",
    "                rs.randint(100, 400), rs.randint(100, 400),\n",
    "                rs.randn() * 2 * np.pi - np.pi\n",
    "                ])\n",
    "        self._set_state(state)\n",
    "\n",
    "        obs = self._get_obs()\n",
    "        info = self._get_info()\n",
    "        return obs, info\n",
    "\n",
    "    def step(self, action):\n",
    "        dt = 1.0 / self.sim_hz\n",
    "        self.n_contact_points = 0\n",
    "        n_steps = self.sim_hz // self.control_hz\n",
    "        if action is not None:\n",
    "            self.latest_action = action\n",
    "            for i in range(n_steps):\n",
    "                # Step PD control.\n",
    "                # self.agent.velocity = self.k_p * (act - self.agent.position)    # P control works too.\n",
    "                acceleration = self.k_p * (action - self.agent.position) + self.k_v * (Vec2d(0, 0) - self.agent.velocity)\n",
    "                self.agent.velocity += acceleration * dt\n",
    "\n",
    "                # Step physics.\n",
    "                self.space.step(dt)\n",
    "\n",
    "        # compute reward\n",
    "        goal_body = self._get_goal_pose_body(self.goal_pose)\n",
    "        goal_geom = pymunk_to_shapely(goal_body, self.block.shapes)\n",
    "        block_geom = pymunk_to_shapely(self.block, self.block.shapes)\n",
    "\n",
    "        intersection_area = goal_geom.intersection(block_geom).area\n",
    "        goal_area = goal_geom.area\n",
    "        coverage = intersection_area / goal_area\n",
    "        reward = np.clip(coverage / self.success_threshold, 0, 1)\n",
    "        done = coverage > self.success_threshold\n",
    "        terminated = done\n",
    "        truncated = done\n",
    "\n",
    "        observation = self._get_obs()\n",
    "        info = self._get_info()\n",
    "\n",
    "        return observation, reward, terminated, truncated, info\n",
    "\n",
    "    def render(self, mode):\n",
    "        return self._render_frame(mode)\n",
    "\n",
    "    def teleop_agent(self):\n",
    "        TeleopAgent = collections.namedtuple('TeleopAgent', ['act'])\n",
    "        def act(obs):\n",
    "            act = None\n",
    "            mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen)\n",
    "            if self.teleop or (mouse_position - self.agent.position).length < 30:\n",
    "                self.teleop = True\n",
    "                act = mouse_position\n",
    "            return act\n",
    "        return TeleopAgent(act)\n",
    "\n",
    "    def _get_obs(self):\n",
    "        obs = np.array(\n",
    "            tuple(self.agent.position) \\\n",
    "            + tuple(self.block.position) \\\n",
    "            + (self.block.angle % (2 * np.pi),))\n",
    "        return obs\n",
    "\n",
    "    def _get_goal_pose_body(self, pose):\n",
    "        mass = 1\n",
    "        inertia = pymunk.moment_for_box(mass, (50, 100))\n",
    "        body = pymunk.Body(mass, inertia)\n",
    "        # preserving the legacy assignment order for compatibility\n",
    "        # the order here dosn't matter somehow, maybe because CoM is aligned with body origin\n",
    "        body.position = pose[:2].tolist()\n",
    "        body.angle = pose[2]\n",
    "        return body\n",
    "\n",
    "    def _get_info(self):\n",
    "        n_steps = self.sim_hz // self.control_hz\n",
    "        n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps))\n",
    "        info = {\n",
    "            'pos_agent': np.array(self.agent.position),\n",
    "            'vel_agent': np.array(self.agent.velocity),\n",
    "            'block_pose': np.array(list(self.block.position) + [self.block.angle]),\n",
    "            'goal_pose': self.goal_pose,\n",
    "            'n_contacts': n_contact_points_per_step}\n",
    "        return info\n",
    "\n",
    "    def _render_frame(self, mode):\n",
    "\n",
    "        if self.window is None and mode == \"human\":\n",
    "            pygame.init()\n",
    "            pygame.display.init()\n",
    "            self.window = pygame.display.set_mode((self.window_size, self.window_size))\n",
    "        if self.clock is None and mode == \"human\":\n",
    "            self.clock = pygame.time.Clock()\n",
    "\n",
    "        canvas = pygame.Surface((self.window_size, self.window_size))\n",
    "        canvas.fill((255, 255, 255))\n",
    "        self.screen = canvas\n",
    "\n",
    "        draw_options = DrawOptions(canvas)\n",
    "\n",
    "        # Draw goal pose.\n",
    "        goal_body = self._get_goal_pose_body(self.goal_pose)\n",
    "        for shape in self.block.shapes:\n",
    "            goal_points = [pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface) for v in shape.get_vertices()]\n",
    "            goal_points += [goal_points[0]]\n",
    "            pygame.draw.polygon(canvas, self.goal_color, goal_points)\n",
    "\n",
    "        # Draw agent and block.\n",
    "        self.space.debug_draw(draw_options)\n",
    "\n",
    "        if mode == \"human\":\n",
    "            # The following line copies our drawings from `canvas` to the visible window\n",
    "            self.window.blit(canvas, canvas.get_rect())\n",
    "            pygame.event.pump()\n",
    "            pygame.display.update()\n",
    "\n",
    "            # the clock is aleady ticked during in step for \"human\"\n",
    "\n",
    "\n",
    "        img = np.transpose(\n",
    "                np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2)\n",
    "            )\n",
    "        img = cv2.resize(img, (self.render_size, self.render_size))\n",
    "        if self.render_action:\n",
    "            if self.render_action and (self.latest_action is not None):\n",
    "                action = np.array(self.latest_action)\n",
    "                coord = (action / 512 * 96).astype(np.int32)\n",
    "                marker_size = int(8/96*self.render_size)\n",
    "                thickness = int(1/96*self.render_size)\n",
    "                cv2.drawMarker(img, coord,\n",
    "                    color=(255,0,0), markerType=cv2.MARKER_CROSS,\n",
    "                    markerSize=marker_size, thickness=thickness)\n",
    "        return img\n",
    "\n",
    "\n",
    "    def close(self):\n",
    "        if self.window is not None:\n",
    "            pygame.display.quit()\n",
    "            pygame.quit()\n",
    "\n",
    "    def seed(self, seed=None):\n",
    "        if seed is None:\n",
    "            seed = np.random.randint(0,25536)\n",
    "        self._seed = seed\n",
    "        self.np_random = np.random.default_rng(seed)\n",
    "\n",
    "    def _handle_collision(self, arbiter, space, data):\n",
    "        self.n_contact_points += len(arbiter.contact_point_set.points)\n",
    "\n",
    "    def _set_state(self, state):\n",
    "        if isinstance(state, np.ndarray):\n",
    "            state = state.tolist()\n",
    "        pos_agent = state[:2]\n",
    "        pos_block = state[2:4]\n",
    "        rot_block = state[4]\n",
    "        self.agent.position = pos_agent\n",
    "        # setting angle rotates with respect to center of mass\n",
    "        # therefore will modify the geometric position\n",
    "        # if not the same as CoM\n",
    "        # therefore should be modified first.\n",
    "        if self.legacy:\n",
    "            # for compatiblity with legacy data\n",
    "            self.block.position = pos_block\n",
    "            self.block.angle = rot_block\n",
    "        else:\n",
    "            self.block.angle = rot_block\n",
    "            self.block.position = pos_block\n",
    "\n",
    "        # Run physics to take effect\n",
    "        self.space.step(1.0 / self.sim_hz)\n",
    "\n",
    "    def _set_state_local(self, state_local):\n",
    "        agent_pos_local = state_local[:2]\n",
    "        block_pose_local = state_local[2:]\n",
    "        tf_img_obj = st.AffineTransform(\n",
    "            translation=self.goal_pose[:2],\n",
    "            rotation=self.goal_pose[2])\n",
    "        tf_obj_new = st.AffineTransform(\n",
    "            translation=block_pose_local[:2],\n",
    "            rotation=block_pose_local[2]\n",
    "        )\n",
    "        tf_img_new = st.AffineTransform(\n",
    "            matrix=tf_img_obj.params @ tf_obj_new.params\n",
    "        )\n",
    "        agent_pos_new = tf_img_new(agent_pos_local)\n",
    "        new_state = np.array(\n",
    "            list(agent_pos_new[0]) + list(tf_img_new.translation) \\\n",
    "                + [tf_img_new.rotation])\n",
    "        self._set_state(new_state)\n",
    "        return new_state\n",
    "\n",
    "    def _setup(self):\n",
    "        self.space = pymunk.Space()\n",
    "        self.space.gravity = 0, 0\n",
    "        self.space.damping = 0\n",
    "        self.teleop = False\n",
    "        self.render_buffer = list()\n",
    "\n",
    "        # Add walls.\n",
    "        walls = [\n",
    "            self._add_segment((5, 506), (5, 5), 2),\n",
    "            self._add_segment((5, 5), (506, 5), 2),\n",
    "            self._add_segment((506, 5), (506, 506), 2),\n",
    "            self._add_segment((5, 506), (506, 506), 2)\n",
    "        ]\n",
    "        self.space.add(*walls)\n",
    "\n",
    "        # Add agent, block, and goal zone.\n",
    "        self.agent = self.add_circle((256, 400), 15)\n",
    "        self.block = self.add_tee((256, 300), 0)\n",
    "        self.goal_color = pygame.Color('LightGreen')\n",
    "        self.goal_pose = np.array([256,256,np.pi/4])  # x, y, theta (in radians)\n",
    "\n",
    "        # Add collision handeling\n",
    "        self.collision_handeler = self.space.add_collision_handler(0, 0)\n",
    "        self.collision_handeler.post_solve = self._handle_collision\n",
    "        self.n_contact_points = 0\n",
    "\n",
    "        self.max_score = 50 * 100\n",
    "        self.success_threshold = 0.95    # 95% coverage.\n",
    "\n",
    "    def _add_segment(self, a, b, radius):\n",
    "        shape = pymunk.Segment(self.space.static_body, a, b, radius)\n",
    "        shape.color = pygame.Color('LightGray')    # https://htmlcolorcodes.com/color-names\n",
    "        return shape\n",
    "\n",
    "    def add_circle(self, position, radius):\n",
    "        body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)\n",
    "        body.position = position\n",
    "        body.friction = 1\n",
    "        shape = pymunk.Circle(body, radius)\n",
    "        shape.color = pygame.Color('RoyalBlue')\n",
    "        self.space.add(body, shape)\n",
    "        return body\n",
    "\n",
    "    def add_box(self, position, height, width):\n",
    "        mass = 1\n",
    "        inertia = pymunk.moment_for_box(mass, (height, width))\n",
    "        body = pymunk.Body(mass, inertia)\n",
    "        body.position = position\n",
    "        shape = pymunk.Poly.create_box(body, (height, width))\n",
    "        shape.color = pygame.Color('LightSlateGray')\n",
    "        self.space.add(body, shape)\n",
    "        return body\n",
    "\n",
    "    def add_tee(self, position, angle, scale=30, color='LightSlateGray', mask=pymunk.ShapeFilter.ALL_MASKS()):\n",
    "        mass = 1\n",
    "        length = 4\n",
    "        vertices1 = [(-length*scale/2, scale),\n",
    "                                 ( length*scale/2, scale),\n",
    "                                 ( length*scale/2, 0),\n",
    "                                 (-length*scale/2, 0)]\n",
    "        inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1)\n",
    "        vertices2 = [(-scale/2, scale),\n",
    "                                 (-scale/2, length*scale),\n",
    "                                 ( scale/2, length*scale),\n",
    "                                 ( scale/2, scale)]\n",
    "        inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1)\n",
    "        body = pymunk.Body(mass, inertia1 + inertia2)\n",
    "        shape1 = pymunk.Poly(body, vertices1)\n",
    "        shape2 = pymunk.Poly(body, vertices2)\n",
    "        shape1.color = pygame.Color(color)\n",
    "        shape2.color = pygame.Color(color)\n",
    "        shape1.filter = pymunk.ShapeFilter(mask=mask)\n",
    "        shape2.filter = pymunk.ShapeFilter(mask=mask)\n",
    "        body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2\n",
    "        body.position = position\n",
    "        body.angle = angle\n",
    "        body.friction = 1\n",
    "        self.space.add(body, shape1, shape2)\n",
    "        return body\n",
    "\n",
    "\n",
    "class PushTImageEnv(PushTEnv):\n",
    "    metadata = {\"render.modes\": [\"rgb_array\"], \"video.frames_per_second\": 10}\n",
    "\n",
    "    def __init__(self,\n",
    "            legacy=False,\n",
    "            block_cog=None,\n",
    "            damping=None,\n",
    "            render_size=96):\n",
    "        super().__init__(\n",
    "            legacy=legacy,\n",
    "            block_cog=block_cog,\n",
    "            damping=damping,\n",
    "            render_size=render_size,\n",
    "            render_action=False)\n",
    "        ws = self.window_size\n",
    "        self.observation_space = spaces.Dict({\n",
    "            'image': spaces.Box(\n",
    "                low=0,\n",
    "                high=1,\n",
    "                shape=(3,render_size,render_size),\n",
    "                dtype=np.float32\n",
    "            ),\n",
    "            'agent_pos': spaces.Box(\n",
    "                low=0,\n",
    "                high=ws,\n",
    "                shape=(2,),\n",
    "                dtype=np.float32\n",
    "            )\n",
    "        })\n",
    "        self.render_cache = None\n",
    "\n",
    "    def _get_obs(self):\n",
    "        img = super()._render_frame(mode='rgb_array')\n",
    "\n",
    "        agent_pos = np.array(self.agent.position)\n",
    "        img_obs = np.moveaxis(img.astype(np.float32) / 255, -1, 0)\n",
    "        obs = {\n",
    "            'image': img_obs,\n",
    "            'agent_pos': agent_pos\n",
    "        }\n",
    "\n",
    "        # draw action\n",
    "        if self.latest_action is not None:\n",
    "            action = np.array(self.latest_action)\n",
    "            coord = (action / 512 * 96).astype(np.int32)\n",
    "            marker_size = int(8/96*self.render_size)\n",
    "            thickness = int(1/96*self.render_size)\n",
    "            cv2.drawMarker(img, coord,\n",
    "                color=(255,0,0), markerType=cv2.MARKER_CROSS,\n",
    "                markerSize=marker_size, thickness=thickness)\n",
    "        self.render_cache = img\n",
    "\n",
    "        return obs\n",
    "\n",
    "    def render(self, mode):\n",
    "        assert mode == 'rgb_array'\n",
    "\n",
    "        if self.render_cache is None:\n",
    "            self._get_obs()\n",
    "\n",
    "        return self.render_cache\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OknH8Qfqrtc9",
    "outputId": "c97c843e-0b23-4c0a-c123-cd72c407e163"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Env Demo**\n",
    "#@markdown Standard Gym Env (0.21.0 API)\n",
    "\n",
    "# 0. create env object\n",
    "env = PushTImageEnv()\n",
    "\n",
    "# 1. seed env for initial state.\n",
    "# Seed 0-200 are used for the demonstration dataset.\n",
    "env.seed(1000)\n",
    "\n",
    "# 2. must reset before use\n",
    "obs, info = env.reset()\n",
    "\n",
    "# 3. 2D positional action space [0,512]\n",
    "action = env.action_space.sample()\n",
    "\n",
    "# 4. Standard gym step method\n",
    "obs, reward, terminated, truncated, info = env.step(action)\n",
    "\n",
    "# prints and explains each dimension of the observation and action vectors\n",
    "with np.printoptions(precision=4, suppress=True, threshold=5):\n",
    "    print(\"obs['image'].shape:\", obs['image'].shape, \"float32, [0,1]\")\n",
    "    print(\"obs['agent_pos'].shape:\", obs['agent_pos'].shape, \"float32, [0,512]\")\n",
    "    print(\"action.shape: \", action.shape, \"float32, [0,512]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "vHepJOFBucwg"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Dataset**\n",
    "#@markdown\n",
    "#@markdown Defines `PushTImageDataset` and helper functions\n",
    "#@markdown\n",
    "#@markdown The dataset class\n",
    "#@markdown - Load data ((image, agent_pos), action) from a zarr storage\n",
    "#@markdown - Normalizes each dimension of agent_pos and action to [-1,1]\n",
    "#@markdown - Returns\n",
    "#@markdown  - All possible segments with length `pred_horizon`\n",
    "#@markdown  - Pads the beginning and the end of each episode with repetition\n",
    "#@markdown  - key `image`: shape (obs_hoirzon, 3, 96, 96)\n",
    "#@markdown  - key `agent_pos`: shape (obs_hoirzon, 2)\n",
    "#@markdown  - key `action`: shape (pred_horizon, 2)\n",
    "\n",
    "def create_sample_indices(\n",
    "        episode_ends:np.ndarray, sequence_length:int,\n",
    "        pad_before: int=0, pad_after: int=0):\n",
    "    indices = list()\n",
    "    for i in range(len(episode_ends)):\n",
    "        start_idx = 0\n",
    "        if i > 0:\n",
    "            start_idx = episode_ends[i-1]\n",
    "        end_idx = episode_ends[i]\n",
    "        episode_length = end_idx - start_idx\n",
    "\n",
    "        min_start = -pad_before\n",
    "        max_start = episode_length - sequence_length + pad_after\n",
    "\n",
    "        # range stops one idx before end\n",
    "        for idx in range(min_start, max_start+1):\n",
    "            buffer_start_idx = max(idx, 0) + start_idx\n",
    "            buffer_end_idx = min(idx+sequence_length, episode_length) + start_idx\n",
    "            start_offset = buffer_start_idx - (idx+start_idx)\n",
    "            end_offset = (idx+sequence_length+start_idx) - buffer_end_idx\n",
    "            sample_start_idx = 0 + start_offset\n",
    "            sample_end_idx = sequence_length - end_offset\n",
    "            indices.append([\n",
    "                buffer_start_idx, buffer_end_idx,\n",
    "                sample_start_idx, sample_end_idx])\n",
    "    indices = np.array(indices)\n",
    "    return indices\n",
    "\n",
    "\n",
    "def sample_sequence(train_data, sequence_length,\n",
    "                    buffer_start_idx, buffer_end_idx,\n",
    "                    sample_start_idx, sample_end_idx):\n",
    "    result = dict()\n",
    "    for key, input_arr in train_data.items():\n",
    "        sample = input_arr[buffer_start_idx:buffer_end_idx]\n",
    "        data = sample\n",
    "        if (sample_start_idx > 0) or (sample_end_idx < sequence_length):\n",
    "            data = np.zeros(\n",
    "                shape=(sequence_length,) + input_arr.shape[1:],\n",
    "                dtype=input_arr.dtype)\n",
    "            if sample_start_idx > 0:\n",
    "                data[:sample_start_idx] = sample[0]\n",
    "            if sample_end_idx < sequence_length:\n",
    "                data[sample_end_idx:] = sample[-1]\n",
    "            data[sample_start_idx:sample_end_idx] = sample\n",
    "        result[key] = data\n",
    "    return result\n",
    "\n",
    "# normalize data\n",
    "def get_data_stats(data):\n",
    "    data = data.reshape(-1,data.shape[-1])\n",
    "    stats = {\n",
    "        'min': np.min(data, axis=0),\n",
    "        'max': np.max(data, axis=0)\n",
    "    }\n",
    "    return stats\n",
    "\n",
    "def normalize_data(data, stats):\n",
    "    # nomalize to [0,1]\n",
    "    ndata = (data - stats['min']) / (stats['max'] - stats['min'])\n",
    "    # normalize to [-1, 1]\n",
    "    ndata = ndata * 2 - 1\n",
    "    return ndata\n",
    "\n",
    "def unnormalize_data(ndata, stats):\n",
    "    ndata = (ndata + 1) / 2\n",
    "    data = ndata * (stats['max'] - stats['min']) + stats['min']\n",
    "    return data\n",
    "\n",
    "# dataset\n",
    "class PushTImageDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self,\n",
    "                 dataset_path: str,\n",
    "                 pred_horizon: int,\n",
    "                 obs_horizon: int,\n",
    "                 action_horizon: int):\n",
    "\n",
    "        # read from zarr dataset\n",
    "        dataset_root = zarr.open(dataset_path, 'r')\n",
    "\n",
    "        # float32, [0,1], (N,96,96,3)\n",
    "        train_image_data = dataset_root['data']['img'][:]\n",
    "        train_image_data = np.moveaxis(train_image_data, -1,1)\n",
    "        # (N,3,96,96)\n",
    "\n",
    "        # (N, D)\n",
    "        train_data = {\n",
    "            # first two dims of state vector are agent (i.e. gripper) locations\n",
    "            'agent_pos': dataset_root['data']['state'][:,:2],\n",
    "            'action': dataset_root['data']['action'][:]\n",
    "        }\n",
    "        episode_ends = dataset_root['meta']['episode_ends'][:]\n",
    "\n",
    "        # compute start and end of each state-action sequence\n",
    "        # also handles padding\n",
    "        indices = create_sample_indices(\n",
    "            episode_ends=episode_ends,\n",
    "            sequence_length=pred_horizon,\n",
    "            pad_before=obs_horizon-1,\n",
    "            pad_after=action_horizon-1)\n",
    "\n",
    "        # compute statistics and normalized data to [-1,1]\n",
    "        stats = dict()\n",
    "        normalized_train_data = dict()\n",
    "        for key, data in train_data.items():\n",
    "            stats[key] = get_data_stats(data)\n",
    "            normalized_train_data[key] = normalize_data(data, stats[key])\n",
    "\n",
    "        # images are already normalized\n",
    "        normalized_train_data['image'] = train_image_data\n",
    "\n",
    "        self.indices = indices\n",
    "        self.stats = stats\n",
    "        self.normalized_train_data = normalized_train_data\n",
    "        self.pred_horizon = pred_horizon\n",
    "        self.action_horizon = action_horizon\n",
    "        self.obs_horizon = obs_horizon\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.indices)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # get the start/end indices for this datapoint\n",
    "        buffer_start_idx, buffer_end_idx, \\\n",
    "            sample_start_idx, sample_end_idx = self.indices[idx]\n",
    "\n",
    "        # get nomralized data using these indices\n",
    "        nsample = sample_sequence(\n",
    "            train_data=self.normalized_train_data,\n",
    "            sequence_length=self.pred_horizon,\n",
    "            buffer_start_idx=buffer_start_idx,\n",
    "            buffer_end_idx=buffer_end_idx,\n",
    "            sample_start_idx=sample_start_idx,\n",
    "            sample_end_idx=sample_end_idx\n",
    "        )\n",
    "\n",
    "        # discard unused observations\n",
    "        nsample['image'] = nsample['image'][:self.obs_horizon,:]\n",
    "        nsample['agent_pos'] = nsample['agent_pos'][:self.obs_horizon,:]\n",
    "        return nsample\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9ZiHF3lzvB6k",
    "outputId": "e6256ff1-ab20-41d4-a9ea-7c035196bf3c"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Dataset Demo**\n",
    "\n",
    "# download demonstration data from Google Drive\n",
    "dataset_path = \"pusht_cchi_v7_replay.zarr.zip\"\n",
    "if not os.path.isfile(dataset_path):\n",
    "    id = \"1KY1InLurpMvJDRb14L9NlXT_fEsCvVUq&confirm=t\"\n",
    "    gdown.download(id=id, output=dataset_path, quiet=False)\n",
    "\n",
    "# parameters\n",
    "pred_horizon = 16\n",
    "obs_horizon = 2\n",
    "action_horizon = 8\n",
    "#|o|o|                             observations: 2\n",
    "#| |a|a|a|a|a|a|a|a|               actions executed: 8\n",
    "#|p|p|p|p|p|p|p|p|p|p|p|p|p|p|p|p| actions predicted: 16\n",
    "\n",
    "# create dataset from file\n",
    "dataset = PushTImageDataset(\n",
    "    dataset_path=dataset_path,\n",
    "    pred_horizon=pred_horizon,\n",
    "    obs_horizon=obs_horizon,\n",
    "    action_horizon=action_horizon\n",
    ")\n",
    "# save training data statistics (min, max) for each dim\n",
    "stats = dataset.stats\n",
    "\n",
    "# create dataloader\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    batch_size=64,\n",
    "    num_workers=4,\n",
    "    shuffle=True,\n",
    "    # accelerate cpu-gpu transfer\n",
    "    pin_memory=True,\n",
    "    # don't kill worker process afte each epoch\n",
    "    persistent_workers=True\n",
    ")\n",
    "\n",
    "# visualize data in batch\n",
    "batch = next(iter(dataloader))\n",
    "print(\"batch['image'].shape:\", batch['image'].shape)\n",
    "print(\"batch['agent_pos'].shape:\", batch['agent_pos'].shape)\n",
    "print(\"batch['action'].shape\", batch['action'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "X-XRB_g3vsgf"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Network**\n",
    "#@markdown\n",
    "#@markdown Defines a 1D UNet architecture `ConditionalUnet1D`\n",
    "#@markdown as the noies prediction network\n",
    "#@markdown\n",
    "#@markdown Components\n",
    "#@markdown - `SinusoidalPosEmb` Positional encoding for the diffusion iteration k\n",
    "#@markdown - `Downsample1d` Strided convolution to reduce temporal resolution\n",
    "#@markdown - `Upsample1d` Transposed convolution to increase temporal resolution\n",
    "#@markdown - `Conv1dBlock` Conv1d --> GroupNorm --> Mish\n",
    "#@markdown - `ConditionalResidualBlock1D` Takes two inputs `x` and `cond`. \\\n",
    "#@markdown `x` is passed through 2 `Conv1dBlock` stacked together with residual connection.\n",
    "#@markdown `cond` is applied to `x` with [FiLM](https://arxiv.org/abs/1709.07871) conditioning.\n",
    "\n",
    "class SinusoidalPosEmb(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, x):\n",
    "        device = x.device\n",
    "        half_dim = self.dim // 2\n",
    "        emb = math.log(10000) / (half_dim - 1)\n",
    "        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)\n",
    "        emb = x[:, None] * emb[None, :]\n",
    "        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n",
    "        return emb\n",
    "\n",
    "\n",
    "class Downsample1d(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv1d(dim, dim, 3, 2, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class Upsample1d(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class Conv1dBlock(nn.Module):\n",
    "    '''\n",
    "        Conv1d --> GroupNorm --> Mish\n",
    "    '''\n",
    "\n",
    "    def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):\n",
    "        super().__init__()\n",
    "\n",
    "        self.block = nn.Sequential(\n",
    "            nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2),\n",
    "            nn.GroupNorm(n_groups, out_channels),\n",
    "            nn.Mish(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.block(x)\n",
    "\n",
    "\n",
    "class ConditionalResidualBlock1D(nn.Module):\n",
    "    def __init__(self,\n",
    "            in_channels,\n",
    "            out_channels,\n",
    "            cond_dim,\n",
    "            kernel_size=3,\n",
    "            n_groups=8):\n",
    "        super().__init__()\n",
    "\n",
    "        self.blocks = nn.ModuleList([\n",
    "            Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups),\n",
    "            Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups),\n",
    "        ])\n",
    "\n",
    "        # FiLM modulation https://arxiv.org/abs/1709.07871\n",
    "        # predicts per-channel scale and bias\n",
    "        cond_channels = out_channels * 2\n",
    "        self.out_channels = out_channels\n",
    "        self.cond_encoder = nn.Sequential(\n",
    "            nn.Mish(),\n",
    "            nn.Linear(cond_dim, cond_channels),\n",
    "            nn.Unflatten(-1, (-1, 1))\n",
    "        )\n",
    "\n",
    "        # make sure dimensions compatible\n",
    "        self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \\\n",
    "            if in_channels != out_channels else nn.Identity()\n",
    "\n",
    "    def forward(self, x, cond):\n",
    "        '''\n",
    "            x : [ batch_size x in_channels x horizon ]\n",
    "            cond : [ batch_size x cond_dim]\n",
    "\n",
    "            returns:\n",
    "            out : [ batch_size x out_channels x horizon ]\n",
    "        '''\n",
    "        out = self.blocks[0](x)\n",
    "        embed = self.cond_encoder(cond)\n",
    "\n",
    "        embed = embed.reshape(\n",
    "            embed.shape[0], 2, self.out_channels, 1)\n",
    "        scale = embed[:,0,...]\n",
    "        bias = embed[:,1,...]\n",
    "        out = scale * out + bias\n",
    "\n",
    "        out = self.blocks[1](out)\n",
    "        out = out + self.residual_conv(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ConditionalUnet1D(nn.Module):\n",
    "    def __init__(self,\n",
    "        input_dim,\n",
    "        global_cond_dim,\n",
    "        diffusion_step_embed_dim=256,\n",
    "        down_dims=[256,512,1024],\n",
    "        kernel_size=5,\n",
    "        n_groups=8\n",
    "        ):\n",
    "        \"\"\"\n",
    "        input_dim: Dim of actions.\n",
    "        global_cond_dim: Dim of global conditioning applied with FiLM\n",
    "          in addition to diffusion step embedding. This is usually obs_horizon * obs_dim\n",
    "        diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k\n",
    "        down_dims: Channel size for each UNet level.\n",
    "          The length of this array determines numebr of levels.\n",
    "        kernel_size: Conv kernel size\n",
    "        n_groups: Number of groups for GroupNorm\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        all_dims = [input_dim] + list(down_dims)\n",
    "        start_dim = down_dims[0]\n",
    "\n",
    "        dsed = diffusion_step_embed_dim\n",
    "        diffusion_step_encoder = nn.Sequential(\n",
    "            SinusoidalPosEmb(dsed),\n",
    "            nn.Linear(dsed, dsed * 4),\n",
    "            nn.Mish(),\n",
    "            nn.Linear(dsed * 4, dsed),\n",
    "        )\n",
    "        cond_dim = dsed + global_cond_dim\n",
    "\n",
    "        in_out = list(zip(all_dims[:-1], all_dims[1:]))\n",
    "        mid_dim = all_dims[-1]\n",
    "        self.mid_modules = nn.ModuleList([\n",
    "            ConditionalResidualBlock1D(\n",
    "                mid_dim, mid_dim, cond_dim=cond_dim,\n",
    "                kernel_size=kernel_size, n_groups=n_groups\n",
    "            ),\n",
    "            ConditionalResidualBlock1D(\n",
    "                mid_dim, mid_dim, cond_dim=cond_dim,\n",
    "                kernel_size=kernel_size, n_groups=n_groups\n",
    "            ),\n",
    "        ])\n",
    "\n",
    "        down_modules = nn.ModuleList([])\n",
    "        for ind, (dim_in, dim_out) in enumerate(in_out):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "            down_modules.append(nn.ModuleList([\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_in, dim_out, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_out, dim_out, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                Downsample1d(dim_out) if not is_last else nn.Identity()\n",
    "            ]))\n",
    "\n",
    "        up_modules = nn.ModuleList([])\n",
    "        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "            up_modules.append(nn.ModuleList([\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_out*2, dim_in, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                ConditionalResidualBlock1D(\n",
    "                    dim_in, dim_in, cond_dim=cond_dim,\n",
    "                    kernel_size=kernel_size, n_groups=n_groups),\n",
    "                Upsample1d(dim_in) if not is_last else nn.Identity()\n",
    "            ]))\n",
    "\n",
    "        final_conv = nn.Sequential(\n",
    "            Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size),\n",
    "            nn.Conv1d(start_dim, input_dim, 1),\n",
    "        )\n",
    "\n",
    "        self.diffusion_step_encoder = diffusion_step_encoder\n",
    "        self.up_modules = up_modules\n",
    "        self.down_modules = down_modules\n",
    "        self.final_conv = final_conv\n",
    "\n",
    "        print(\"number of parameters: {:e}\".format(\n",
    "            sum(p.numel() for p in self.parameters()))\n",
    "        )\n",
    "\n",
    "    def forward(self,\n",
    "            sample: torch.Tensor,\n",
    "            timestep: Union[torch.Tensor, float, int],\n",
    "            global_cond=None):\n",
    "        \"\"\"\n",
    "        x: (B,T,input_dim)\n",
    "        timestep: (B,) or int, diffusion step\n",
    "        global_cond: (B,global_cond_dim)\n",
    "        output: (B,T,input_dim)\n",
    "        \"\"\"\n",
    "        # (B,T,C)\n",
    "        sample = sample.moveaxis(-1,-2)\n",
    "        # (B,C,T)\n",
    "\n",
    "        # 1. time\n",
    "        timesteps = timestep\n",
    "        if not torch.is_tensor(timesteps):\n",
    "            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can\n",
    "            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)\n",
    "        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:\n",
    "            timesteps = timesteps[None].to(sample.device)\n",
    "        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n",
    "        timesteps = timesteps.expand(sample.shape[0])\n",
    "\n",
    "        global_feature = self.diffusion_step_encoder(timesteps)\n",
    "\n",
    "        if global_cond is not None:\n",
    "            global_feature = torch.cat([\n",
    "                global_feature, global_cond\n",
    "            ], axis=-1)\n",
    "\n",
    "        x = sample\n",
    "        h = []\n",
    "        for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules):\n",
    "            x = resnet(x, global_feature)\n",
    "            x = resnet2(x, global_feature)\n",
    "            h.append(x)\n",
    "            x = downsample(x)\n",
    "\n",
    "        for mid_module in self.mid_modules:\n",
    "            x = mid_module(x, global_feature)\n",
    "\n",
    "        for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules):\n",
    "            x = torch.cat((x, h.pop()), dim=1)\n",
    "            x = resnet(x, global_feature)\n",
    "            x = resnet2(x, global_feature)\n",
    "            x = upsample(x)\n",
    "\n",
    "        x = self.final_conv(x)\n",
    "\n",
    "        # (B,C,T)\n",
    "        x = x.moveaxis(-1,-2)\n",
    "        # (B,T,C)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yXq4r744aMh1"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Vision Encoder**\n",
    "#@markdown\n",
    "#@markdown Defines helper functions:\n",
    "#@markdown - `get_resnet` to initialize standard ResNet vision encoder\n",
    "#@markdown - `replace_bn_with_gn` to replace all BatchNorm layers with GroupNorm\n",
    "\n",
    "def get_resnet(name:str, weights=None, **kwargs) -> nn.Module:\n",
    "    \"\"\"\n",
    "    name: resnet18, resnet34, resnet50\n",
    "    weights: \"IMAGENET1K_V1\", None\n",
    "    \"\"\"\n",
    "    # Use standard ResNet implementation from torchvision\n",
    "    func = getattr(torchvision.models, name)\n",
    "    resnet = func(weights=weights, **kwargs)\n",
    "\n",
    "    # remove the final fully connected layer\n",
    "    # for resnet18, the output dim should be 512\n",
    "    resnet.fc = torch.nn.Identity()\n",
    "    return resnet\n",
    "\n",
    "\n",
    "def replace_submodules(\n",
    "        root_module: nn.Module,\n",
    "        predicate: Callable[[nn.Module], bool],\n",
    "        func: Callable[[nn.Module], nn.Module]) -> nn.Module:\n",
    "    \"\"\"\n",
    "    Replace all submodules selected by the predicate with\n",
    "    the output of func.\n",
    "\n",
    "    predicate: Return true if the module is to be replaced.\n",
    "    func: Return new module to use.\n",
    "    \"\"\"\n",
    "    if predicate(root_module):\n",
    "        return func(root_module)\n",
    "\n",
    "    bn_list = [k.split('.') for k, m\n",
    "        in root_module.named_modules(remove_duplicate=True)\n",
    "        if predicate(m)]\n",
    "    for *parent, k in bn_list:\n",
    "        parent_module = root_module\n",
    "        if len(parent) > 0:\n",
    "            parent_module = root_module.get_submodule('.'.join(parent))\n",
    "        if isinstance(parent_module, nn.Sequential):\n",
    "            src_module = parent_module[int(k)]\n",
    "        else:\n",
    "            src_module = getattr(parent_module, k)\n",
    "        tgt_module = func(src_module)\n",
    "        if isinstance(parent_module, nn.Sequential):\n",
    "            parent_module[int(k)] = tgt_module\n",
    "        else:\n",
    "            setattr(parent_module, k, tgt_module)\n",
    "    # verify that all modules are replaced\n",
    "    bn_list = [k.split('.') for k, m\n",
    "        in root_module.named_modules(remove_duplicate=True)\n",
    "        if predicate(m)]\n",
    "    assert len(bn_list) == 0\n",
    "    return root_module\n",
    "\n",
    "def replace_bn_with_gn(\n",
    "    root_module: nn.Module,\n",
    "    features_per_group: int=16) -> nn.Module:\n",
    "    \"\"\"\n",
    "    Relace all BatchNorm layers with GroupNorm.\n",
    "    \"\"\"\n",
    "    replace_submodules(\n",
    "        root_module=root_module,\n",
    "        predicate=lambda x: isinstance(x, nn.BatchNorm2d),\n",
    "        func=lambda x: nn.GroupNorm(\n",
    "            num_groups=x.num_features//features_per_group,\n",
    "            num_channels=x.num_features)\n",
    "    )\n",
    "    return root_module\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4APZkqh336-M",
    "outputId": "e362d582-ceac-4cdd-e866-c34858593696"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Network Demo**\n",
    "\n",
    "# construct ResNet18 encoder\n",
    "# if you have multiple camera views, use seperate encoder weights for each view.\n",
    "vision_encoder = get_resnet('resnet18')\n",
    "\n",
    "# IMPORTANT!\n",
    "# replace all BatchNorm with GroupNorm to work with EMA\n",
    "# performance will tank if you forget to do this!\n",
    "vision_encoder = replace_bn_with_gn(vision_encoder)\n",
    "\n",
    "# ResNet18 has output dim of 512\n",
    "vision_feature_dim = 512\n",
    "# agent_pos is 2 dimensional\n",
    "lowdim_obs_dim = 2\n",
    "# observation feature has 514 dims in total per step\n",
    "obs_dim = vision_feature_dim + lowdim_obs_dim\n",
    "action_dim = 2\n",
    "\n",
    "# create network object\n",
    "noise_pred_net = ConditionalUnet1D(\n",
    "    input_dim=action_dim,\n",
    "    global_cond_dim=obs_dim*obs_horizon\n",
    ")\n",
    "\n",
    "# the final arch has 2 parts\n",
    "nets = nn.ModuleDict({\n",
    "    'vision_encoder': vision_encoder,\n",
    "    'noise_pred_net': noise_pred_net\n",
    "})\n",
    "\n",
    "# demo\n",
    "with torch.no_grad():\n",
    "    # example inputs\n",
    "    image = torch.zeros((1, obs_horizon,3,96,96))\n",
    "    agent_pos = torch.zeros((1, obs_horizon, 2))\n",
    "    # vision encoder\n",
    "    image_features = nets['vision_encoder'](\n",
    "        image.flatten(end_dim=1))\n",
    "    # (2,512)\n",
    "    image_features = image_features.reshape(*image.shape[:2],-1)\n",
    "    # (1,2,512)\n",
    "    obs = torch.cat([image_features, agent_pos],dim=-1)\n",
    "    # (1,2,514)\n",
    "\n",
    "    noised_action = torch.randn((1, pred_horizon, action_dim))\n",
    "    diffusion_iter = torch.zeros((1,))\n",
    "\n",
    "    # the noise prediction network\n",
    "    # takes noisy action, diffusion iteration and observation as input\n",
    "    # predicts the noise added to action\n",
    "    noise = nets['noise_pred_net'](\n",
    "        sample=noised_action,\n",
    "        timestep=diffusion_iter,\n",
    "        global_cond=obs.flatten(start_dim=1))\n",
    "\n",
    "    # illustration of removing noise\n",
    "    # the actual noise removal is performed by NoiseScheduler\n",
    "    # and is dependent on the diffusion noise schedule\n",
    "    denoised_action = noised_action - noise\n",
    "\n",
    "# for this demo, we use DDPMScheduler with 100 diffusion iterations\n",
    "num_diffusion_iters = 100\n",
    "noise_scheduler = DDPMScheduler(\n",
    "    num_train_timesteps=num_diffusion_iters,\n",
    "    # the choise of beta schedule has big impact on performance\n",
    "    # we found squared cosine works the best\n",
    "    beta_schedule='squaredcos_cap_v2',\n",
    "    # clip output to [-1,1] to improve stability\n",
    "    clip_sample=True,\n",
    "    # our network predicts noise (instead of denoised action)\n",
    "    prediction_type='epsilon'\n",
    ")\n",
    "\n",
    "# device transfer\n",
    "device = torch.device('cuda')\n",
    "_ = nets.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 81,
     "referenced_widgets": [
      "b34bf5a5d603446fb453bbed31c4469c",
      "0873a2dee0e44b0fb3e8445d94c27171",
      "410476fa4ac14dd2b78377f4d779402f",
      "09152dbe3c8543aa809aa592afdc53f1",
      "e2a91fa6d1514913892c0def2e2ecf78",
      "7fd765bcd2a34db49a02839ba1c4f518",
      "d9054c34284045eda3546a21bb5fafc3",
      "a92fc95c6a3d496997da90c87e24ba01",
      "9a98272475d940b1b98438c4d02dcc05",
      "47b6383ed3b64b7e8c7e66d0e7a468d9",
      "1cdde5afac8041bf9b75e0e80bd80630",
      "70b4b657d68648d9be98c099d061cd5a",
      "8c0237e38cfd4143aa423dea26637965",
      "44c486fa8c4241f4a1a245f4e24da769",
      "12ef3c026774405cb91faf18d7fd8afa",
      "99c0577d549c4dd4a61d9ec0cf6254d5",
      "5124950790874225b5cca7556f122f9a",
      "f609505b750747dcaff8a950131743e9",
      "3c0a27b6addb4c6a977824995548fe90",
      "d14f2c65ef9348348250af2df0f32963",
      "eb23edb22d624aee9fd55ab05e2a517e",
      "a073d659b0f347c9b30c343430d7f6aa",
      "7508fbad995648a39d2a55953d000786",
      "9009dc5a656c4befbf513f2574df5a7d",
      "479d2f35498b4fbea77057f3f26fd287",
      "16b5928585984a1b81b809717f125489",
      "9d436c7ee90349a5966f025d095dd0cc",
      "3e503eeff5d94f0e9dc8a3d31190c6a2",
      "4abc5a8b257240b99535a6abe5e00ef6",
      "6c25264289e3411db29f6f8db936b00f",
      "a0ff8cd1b00544caa12cbd809524dd15",
      "1faac530b70b41a58a8c9c0cf44af69c",
      "7cfefcf49a9a426da4e9203ec4debe38",
      "5f96241543b441fcb13c16d4ca476405",
      "e6bbb6ce052045ca895b526c620d8847",
      "5c1c61cd14fb4aa8b335684fbf652a29",
      "a5a2bc5dc0b54a4796e574308bb6c4bc",
      "febe66177c604395a78be993eeff7c6d",
      "6de3445a866442aa9cc7d855cb0d0740",
      "353cb60a2a73417b81ccbb5e52480ed3",
      "69f323a406004822bf1a0bfd79767c9a",
      "88a8f5924a1e4a86805e314dd90c17be",
      "9ada44d21f234712ac5df797c730495a",
      "c777ac64366c40ea8a1f63408879c2ab"
     ]
    },
    "id": "93E9RdnR4D8v",
    "outputId": "6592f2ca-4d83-41a4-c650-564042a44b0d"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Training**\n",
    "#@markdown\n",
    "#@markdown Takes about 2.5 hours. If you don't want to wait, skip to the next cell\n",
    "#@markdown to load pre-trained weights\n",
    "\n",
    "num_epochs = 100\n",
    "\n",
    "# Exponential Moving Average\n",
    "# accelerates training and improves stability\n",
    "# holds a copy of the model weights\n",
    "ema = EMAModel(\n",
    "    parameters=nets.parameters(),\n",
    "    power=0.75)\n",
    "\n",
    "# Standard ADAM optimizer\n",
    "# Note that EMA parametesr are not optimized\n",
    "optimizer = torch.optim.AdamW(\n",
    "    params=nets.parameters(),\n",
    "    lr=1e-4, weight_decay=1e-6)\n",
    "\n",
    "# Cosine LR schedule with linear warmup\n",
    "lr_scheduler = get_scheduler(\n",
    "    name='cosine',\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=500,\n",
    "    num_training_steps=len(dataloader) * num_epochs\n",
    ")\n",
    "\n",
    "with tqdm(range(num_epochs), desc='Epoch') as tglobal:\n",
    "    # epoch loop\n",
    "    for epoch_idx in tglobal:\n",
    "        epoch_loss = list()\n",
    "        # batch loop\n",
    "        with tqdm(dataloader, desc='Batch', leave=False) as tepoch:\n",
    "            for nbatch in tepoch:\n",
    "                # data normalized in dataset\n",
    "                # device transfer\n",
    "                nimage = nbatch['image'][:,:obs_horizon].to(device)\n",
    "                nagent_pos = nbatch['agent_pos'][:,:obs_horizon].to(device)\n",
    "                naction = nbatch['action'].to(device)\n",
    "                B = nagent_pos.shape[0]\n",
    "\n",
    "                # encoder vision features\n",
    "                image_features = nets['vision_encoder'](\n",
    "                    nimage.flatten(end_dim=1))\n",
    "                image_features = image_features.reshape(\n",
    "                    *nimage.shape[:2],-1)\n",
    "                # (B,obs_horizon,D)\n",
    "\n",
    "                # concatenate vision feature and low-dim obs\n",
    "                obs_features = torch.cat([image_features, nagent_pos], dim=-1)\n",
    "                obs_cond = obs_features.flatten(start_dim=1)\n",
    "                # (B, obs_horizon * obs_dim)\n",
    "\n",
    "                # sample noise to add to actions\n",
    "                noise = torch.randn(naction.shape, device=device)\n",
    "\n",
    "                # sample a diffusion iteration for each data point\n",
    "                timesteps = torch.randint(\n",
    "                    0, noise_scheduler.config.num_train_timesteps,\n",
    "                    (B,), device=device\n",
    "                ).long()\n",
    "\n",
    "                # add noise to the clean images according to the noise magnitude at each diffusion iteration\n",
    "                # (this is the forward diffusion process)\n",
    "                noisy_actions = noise_scheduler.add_noise(\n",
    "                    naction, noise, timesteps)\n",
    "\n",
    "                # predict the noise residual\n",
    "                noise_pred = noise_pred_net(\n",
    "                    noisy_actions, timesteps, global_cond=obs_cond)\n",
    "\n",
    "                # L2 loss\n",
    "                loss = nn.functional.mse_loss(noise_pred, noise)\n",
    "\n",
    "                # optimize\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "                # step lr scheduler every batch\n",
    "                # this is different from standard pytorch behavior\n",
    "                lr_scheduler.step()\n",
    "\n",
    "                # update Exponential Moving Average of the model weights\n",
    "                ema.step(nets.parameters())\n",
    "\n",
    "                # logging\n",
    "                loss_cpu = loss.item()\n",
    "                epoch_loss.append(loss_cpu)\n",
    "                tepoch.set_postfix(loss=loss_cpu)\n",
    "        tglobal.set_postfix(loss=np.mean(epoch_loss))\n",
    "\n",
    "# Weights of the EMA model\n",
    "# is used for inference\n",
    "ema_nets = nets\n",
    "ema.copy_to(ema_nets.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6F3hUbIuxGdO",
    "outputId": "36fa4685-4c4a-4ef3-a0a5-add134288f88"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Loading Pretrained Checkpoint**\n",
    "#@markdown Set `load_pretrained = True` to load pretrained weights.\n",
    "\n",
    "load_pretrained = False\n",
    "if load_pretrained:\n",
    "  ckpt_path = \"pusht_vision_100ep.ckpt\"\n",
    "  if not os.path.isfile(ckpt_path):\n",
    "      id = \"1XKpfNSlwYMGaF5CncoFaLKCDTWoLAHf1&confirm=t\"\n",
    "      gdown.download(id=id, output=ckpt_path, quiet=False)\n",
    "\n",
    "  state_dict = torch.load(ckpt_path, map_location='cuda')\n",
    "  ema_nets = nets\n",
    "  ema_nets.load_state_dict(state_dict)\n",
    "  print('Pretrained weights loaded.')\n",
    "else:\n",
    "  print(\"Skipped pretrained weight loading.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 327,
     "referenced_widgets": [
      "b0a9f64f47d34769a47e165291d2c550",
      "2c31ffecd0494e2a8d0a8cefbea5df8a",
      "1d0a57ed8e914d31bcc1c36a56451df7",
      "e1bb86eb510b47d9ac2311926455bf1f",
      "182afbf676cd4c0982b90890bbdbeeef",
      "5d4a78f691ec4555bf179e5ef5d828ac",
      "1b96a78d8b8542eb830e3c446a0dcf42",
      "4cc800669dfd45a784864dd4a3881daa",
      "95dc6c44f97e4e8882a3657bd2fd66fb",
      "333b5148e59e47f3998a5cf5df531baf",
      "185ce2207cec4ae1a6ffdecadae23894"
     ]
    },
    "id": "OyLjlNQk5nr9",
    "outputId": "4641c055-a534-48e2-cc26-0d1056a004ba"
   },
   "outputs": [],
   "source": [
    "#@markdown ### **Inference**\n",
    "\n",
    "# limit enviornment interaction to 200 steps before termination\n",
    "max_steps = 200\n",
    "env = PushTImageEnv()\n",
    "# use a seed >200 to avoid initial states seen in the training dataset\n",
    "env.seed(100000)\n",
    "\n",
    "# get first observation\n",
    "obs, info = env.reset()\n",
    "\n",
    "# keep a queue of last 2 steps of observations\n",
    "obs_deque = collections.deque(\n",
    "    [obs] * obs_horizon, maxlen=obs_horizon)\n",
    "# save visualization and rewards\n",
    "imgs = [env.render(mode='rgb_array')]\n",
    "rewards = list()\n",
    "done = False\n",
    "step_idx = 0\n",
    "\n",
    "with tqdm(total=max_steps, desc=\"Eval PushTImageEnv\") as pbar:\n",
    "    while not done:\n",
    "        B = 1\n",
    "        # stack the last obs_horizon number of observations\n",
    "        images = np.stack([x['image'] for x in obs_deque])\n",
    "        agent_poses = np.stack([x['agent_pos'] for x in obs_deque])\n",
    "\n",
    "        # normalize observation\n",
    "        nagent_poses = normalize_data(agent_poses, stats=stats['agent_pos'])\n",
    "        # images are already normalized to [0,1]\n",
    "        nimages = images\n",
    "\n",
    "        # device transfer\n",
    "        nimages = torch.from_numpy(nimages).to(device, dtype=torch.float32)\n",
    "        # (2,3,96,96)\n",
    "        nagent_poses = torch.from_numpy(nagent_poses).to(device, dtype=torch.float32)\n",
    "        # (2,2)\n",
    "\n",
    "        # infer action\n",
    "        with torch.no_grad():\n",
    "            # get image features\n",
    "            image_features = ema_nets['vision_encoder'](nimages)\n",
    "            # (2,512)\n",
    "\n",
    "            # concat with low-dim observations\n",
    "            obs_features = torch.cat([image_features, nagent_poses], dim=-1)\n",
    "\n",
    "            # reshape observation to (B,obs_horizon*obs_dim)\n",
    "            obs_cond = obs_features.unsqueeze(0).flatten(start_dim=1)\n",
    "\n",
    "            # initialize action from Guassian noise\n",
    "            noisy_action = torch.randn(\n",
    "                (B, pred_horizon, action_dim), device=device)\n",
    "            naction = noisy_action\n",
    "\n",
    "            # init scheduler\n",
    "            noise_scheduler.set_timesteps(num_diffusion_iters)\n",
    "\n",
    "            for k in noise_scheduler.timesteps:\n",
    "                # predict noise\n",
    "                noise_pred = ema_nets['noise_pred_net'](\n",
    "                    sample=naction,\n",
    "                    timestep=k,\n",
    "                    global_cond=obs_cond\n",
    "                )\n",
    "\n",
    "                # inverse diffusion step (remove noise)\n",
    "                naction = noise_scheduler.step(\n",
    "                    model_output=noise_pred,\n",
    "                    timestep=k,\n",
    "                    sample=naction\n",
    "                ).prev_sample\n",
    "\n",
    "        # unnormalize action\n",
    "        naction = naction.detach().to('cpu').numpy()\n",
    "        # (B, pred_horizon, action_dim)\n",
    "        naction = naction[0]\n",
    "        action_pred = unnormalize_data(naction, stats=stats['action'])\n",
    "\n",
    "        # only take action_horizon number of actions\n",
    "        start = obs_horizon - 1\n",
    "        end = start + action_horizon\n",
    "        action = action_pred[start:end,:]\n",
    "        # (action_horizon, action_dim)\n",
    "\n",
    "        # execute action_horizon number of steps\n",
    "        # without replanning\n",
    "        for i in range(len(action)):\n",
    "            # stepping env\n",
    "            obs, reward, done, _, info = env.step(action[i])\n",
    "            # save observations\n",
    "            obs_deque.append(obs)\n",
    "            # and reward/vis\n",
    "            rewards.append(reward)\n",
    "            imgs.append(env.render(mode='rgb_array'))\n",
    "\n",
    "            # update progress bar\n",
    "            step_idx += 1\n",
    "            pbar.update(1)\n",
    "            pbar.set_postfix(reward=reward)\n",
    "            if step_idx > max_steps:\n",
    "                done = True\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "# print out the maximum target coverage\n",
    "print('Score: ', max(rewards))\n",
    "\n",
    "# visualize\n",
    "from IPython.display import Video\n",
    "vwrite('vis.mp4', imgs)\n",
    "Video('vis.mp4', embed=True, width=256, height=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0873a2dee0e44b0fb3e8445d94c27171": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7fd765bcd2a34db49a02839ba1c4f518",
      "placeholder": "​",
      "style": "IPY_MODEL_d9054c34284045eda3546a21bb5fafc3",
      "value": "Epoch:   2%"
     }
    },
    "09152dbe3c8543aa809aa592afdc53f1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_47b6383ed3b64b7e8c7e66d0e7a468d9",
      "placeholder": "​",
      "style": "IPY_MODEL_1cdde5afac8041bf9b75e0e80bd80630",
      "value": " 2/100 [03:14&lt;2:38:57, 97.32s/it, loss=0.0192]"
     }
    },
    "11d14b9bf0f74bf7bd1b58e6a26fd49e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "12ef3c026774405cb91faf18d7fd8afa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_eb23edb22d624aee9fd55ab05e2a517e",
      "placeholder": "​",
      "style": "IPY_MODEL_a073d659b0f347c9b30c343430d7f6aa",
      "value": " 379/379 [01:37&lt;00:00,  4.40it/s, loss=0.0185]"
     }
    },
    "16b5928585984a1b81b809717f125489": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1faac530b70b41a58a8c9c0cf44af69c",
      "placeholder": "​",
      "style": "IPY_MODEL_7cfefcf49a9a426da4e9203ec4debe38",
      "value": " 379/379 [01:37&lt;00:00,  4.57it/s, loss=0.0103]"
     }
    },
    "182afbf676cd4c0982b90890bbdbeeef": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "185ce2207cec4ae1a6ffdecadae23894": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "19e913a07b044b2d91ca7b5d2dbbe8f7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_11d14b9bf0f74bf7bd1b58e6a26fd49e",
      "placeholder": "​",
      "style": "IPY_MODEL_e9fde27bd58b499bb386b36314d9efca",
      "value": ""
     }
    },
    "1b96a78d8b8542eb830e3c446a0dcf42": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1cdde5afac8041bf9b75e0e80bd80630": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1d0a57ed8e914d31bcc1c36a56451df7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4cc800669dfd45a784864dd4a3881daa",
      "max": 200,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_95dc6c44f97e4e8882a3657bd2fd66fb",
      "value": 200
     }
    },
    "1de97479fc3e44caae336cf0b1082d02": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5407db80343c44c79350021cd98204f8",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_68e9f5994766480dac5dc77ce043a2e9",
      "value": 0
     }
    },
    "1faac530b70b41a58a8c9c0cf44af69c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "2c31ffecd0494e2a8d0a8cefbea5df8a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5d4a78f691ec4555bf179e5ef5d828ac",
      "placeholder": "​",
      "style": "IPY_MODEL_1b96a78d8b8542eb830e3c446a0dcf42",
      "value": "Eval PushTImageEnv: "
     }
    },
    "333b5148e59e47f3998a5cf5df531baf": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "353cb60a2a73417b81ccbb5e52480ed3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "3c0a27b6addb4c6a977824995548fe90": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3e503eeff5d94f0e9dc8a3d31190c6a2": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "410476fa4ac14dd2b78377f4d779402f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a92fc95c6a3d496997da90c87e24ba01",
      "max": 100,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_9a98272475d940b1b98438c4d02dcc05",
      "value": 2
     }
    },
    "44c486fa8c4241f4a1a245f4e24da769": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3c0a27b6addb4c6a977824995548fe90",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_d14f2c65ef9348348250af2df0f32963",
      "value": 379
     }
    },
    "479d2f35498b4fbea77057f3f26fd287": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6c25264289e3411db29f6f8db936b00f",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_a0ff8cd1b00544caa12cbd809524dd15",
      "value": 379
     }
    },
    "47b6383ed3b64b7e8c7e66d0e7a468d9": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4abc5a8b257240b99535a6abe5e00ef6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "4cc800669dfd45a784864dd4a3881daa": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5124950790874225b5cca7556f122f9a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5407db80343c44c79350021cd98204f8": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "54c35b52c4e24543872c8335c935c410": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "5c1c61cd14fb4aa8b335684fbf652a29": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_69f323a406004822bf1a0bfd79767c9a",
      "max": 379,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_88a8f5924a1e4a86805e314dd90c17be",
      "value": 11
     }
    },
    "5d4a78f691ec4555bf179e5ef5d828ac": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5f96241543b441fcb13c16d4ca476405": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_e6bbb6ce052045ca895b526c620d8847",
       "IPY_MODEL_5c1c61cd14fb4aa8b335684fbf652a29",
       "IPY_MODEL_a5a2bc5dc0b54a4796e574308bb6c4bc"
      ],
      "layout": "IPY_MODEL_febe66177c604395a78be993eeff7c6d"
     }
    },
    "68e9f5994766480dac5dc77ce043a2e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "69f323a406004822bf1a0bfd79767c9a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6c25264289e3411db29f6f8db936b00f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6de3445a866442aa9cc7d855cb0d0740": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "70b4b657d68648d9be98c099d061cd5a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_8c0237e38cfd4143aa423dea26637965",
       "IPY_MODEL_44c486fa8c4241f4a1a245f4e24da769",
       "IPY_MODEL_12ef3c026774405cb91faf18d7fd8afa"
      ],
      "layout": "IPY_MODEL_99c0577d549c4dd4a61d9ec0cf6254d5"
     }
    },
    "7508fbad995648a39d2a55953d000786": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_9009dc5a656c4befbf513f2574df5a7d",
       "IPY_MODEL_479d2f35498b4fbea77057f3f26fd287",
       "IPY_MODEL_16b5928585984a1b81b809717f125489"
      ],
      "layout": "IPY_MODEL_9d436c7ee90349a5966f025d095dd0cc"
     }
    },
    "7cfefcf49a9a426da4e9203ec4debe38": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "7fd765bcd2a34db49a02839ba1c4f518": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "88a8f5924a1e4a86805e314dd90c17be": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8c0237e38cfd4143aa423dea26637965": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5124950790874225b5cca7556f122f9a",
      "placeholder": "​",
      "style": "IPY_MODEL_f609505b750747dcaff8a950131743e9",
      "value": "Batch: 100%"
     }
    },
    "9009dc5a656c4befbf513f2574df5a7d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3e503eeff5d94f0e9dc8a3d31190c6a2",
      "placeholder": "​",
      "style": "IPY_MODEL_4abc5a8b257240b99535a6abe5e00ef6",
      "value": "Batch: 100%"
     }
    },
    "95dc6c44f97e4e8882a3657bd2fd66fb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "99c0577d549c4dd4a61d9ec0cf6254d5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "9a98272475d940b1b98438c4d02dcc05": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "9ada44d21f234712ac5df797c730495a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9d436c7ee90349a5966f025d095dd0cc": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "a073d659b0f347c9b30c343430d7f6aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "a0ff8cd1b00544caa12cbd809524dd15": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "a5a2bc5dc0b54a4796e574308bb6c4bc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9ada44d21f234712ac5df797c730495a",
      "placeholder": "​",
      "style": "IPY_MODEL_c777ac64366c40ea8a1f63408879c2ab",
      "value": " 11/379 [00:03&lt;01:34,  3.91it/s, loss=0.0256]"
     }
    },
    "a8e97e26d65147efa2ddb176214e818a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_19e913a07b044b2d91ca7b5d2dbbe8f7",
       "IPY_MODEL_1de97479fc3e44caae336cf0b1082d02",
       "IPY_MODEL_b8d4e152c232416f82f5fb30e2478ef8"
      ],
      "layout": "IPY_MODEL_d66b3f8971b24c81abaefe111e6383a6"
     }
    },
    "a92fc95c6a3d496997da90c87e24ba01": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b0a9f64f47d34769a47e165291d2c550": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_2c31ffecd0494e2a8d0a8cefbea5df8a",
       "IPY_MODEL_1d0a57ed8e914d31bcc1c36a56451df7",
       "IPY_MODEL_e1bb86eb510b47d9ac2311926455bf1f"
      ],
      "layout": "IPY_MODEL_182afbf676cd4c0982b90890bbdbeeef"
     }
    },
    "b34bf5a5d603446fb453bbed31c4469c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_0873a2dee0e44b0fb3e8445d94c27171",
       "IPY_MODEL_410476fa4ac14dd2b78377f4d779402f",
       "IPY_MODEL_09152dbe3c8543aa809aa592afdc53f1"
      ],
      "layout": "IPY_MODEL_e2a91fa6d1514913892c0def2e2ecf78"
     }
    },
    "b8d4e152c232416f82f5fb30e2478ef8": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d00c195ac5644bd0aba81dccf56e2df0",
      "placeholder": "​",
      "style": "IPY_MODEL_54c35b52c4e24543872c8335c935c410",
      "value": " 0/0 [00:00&lt;?, ?it/s]"
     }
    },
    "c777ac64366c40ea8a1f63408879c2ab": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d00c195ac5644bd0aba81dccf56e2df0": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d14f2c65ef9348348250af2df0f32963": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "d66b3f8971b24c81abaefe111e6383a6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d9054c34284045eda3546a21bb5fafc3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e1bb86eb510b47d9ac2311926455bf1f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_333b5148e59e47f3998a5cf5df531baf",
      "placeholder": "​",
      "style": "IPY_MODEL_185ce2207cec4ae1a6ffdecadae23894",
      "value": " 201/? [00:34&lt;00:00,  6.11it/s, reward=0.828]"
     }
    },
    "e2a91fa6d1514913892c0def2e2ecf78": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e6bbb6ce052045ca895b526c620d8847": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6de3445a866442aa9cc7d855cb0d0740",
      "placeholder": "​",
      "style": "IPY_MODEL_353cb60a2a73417b81ccbb5e52480ed3",
      "value": "Batch:   3%"
     }
    },
    "e9fde27bd58b499bb386b36314d9efca": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "eb23edb22d624aee9fd55ab05e2a517e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "f609505b750747dcaff8a950131743e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "febe66177c604395a78be993eeff7c6d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
